{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "introduction",
   "metadata": {},
   "source": [
    "# Efficient Llama Training with Gradient Checkpointing and _Adapters_\n",
    "\n",
    "In this notebook, we show how to efficiently fine-tune a **Llama 3** model using **gradient checkpointing** and adapter methods.\n",
    "\n",
    "**Gradient checkpointing** is a technique to reduce peak memory usage significantly and thus enables training larger models with larger batch sizes. Gradient checkpointing achieves this by trading compute for memory: During the forward pass, gradient checkpointing only stores a subset of activations (thus saving memory). During backpropagation, gradient checkpointing recomputes the activations that were not stored. This can significantly reduce memory requirements at the cost of slightly increased computation time.\n",
    "\n",
    "In this notebook, we finetune Llama-3 8B on supervised instruction tuning data collected by the [Open Assistant project](https://github.com/LAION-AI/Open-Assistant) for training chatbots.\n",
    "\n",
    "Another way to reduce memore usage is to use quantization. Have a look a the [QLora notebook](QLoRA_Llama_Finetuning.ipynb) for an example. This gradient checkpointing notebook is based on the QLoRA notebook. While we use a normal LoRA setup in this notebook, you can easily replace LoRA with QLoRA to reduce memory usage even further."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "installation",
   "metadata": {},
   "source": [
    "## Installation\n",
    "\n",
    "We need `adapters`, `datasets` and `pytorch` for training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "install",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -qq -U adapters datasets torch"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dataset",
   "metadata": {},
   "source": [
    "## Load Open Assistant dataset\n",
    "\n",
    "We use the [`timdettmers/openassistant-guanaco`](https://huggingface.co/datasets/timdettmers/openassistant-guanaco) dataset, which contains a small subset of conversations from the full Open Assistant database."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "load_dataset",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['text'],\n",
       "        num_rows: 9846\n",
       "    })\n",
       "    test: Dataset({\n",
       "        features: ['text'],\n",
       "        num_rows: 518\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "dataset = load_dataset(\"timdettmers/openassistant-guanaco\")\n",
    "dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "model_setup",
   "metadata": {},
   "source": [
    "## Load and prepare model\n",
    "\n",
    "We download the official Llama-2 7B/ Llama-3 8B checkpoint from the HuggingFace Hub. Note that you must request access to this model on the HuggingFace website and use an API token to download it.\n",
    "\n",
    "The key difference in this notebook is that we'll enable gradient checkpointing to reduce memory usage during training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "load_model",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "83e60dee3c434bb3a2bc656bd7f4b667",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/4 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import torch\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "\n",
    "# To access the Llama 3 model, you need to provide your key:\n",
    "HUGGINGFACE_ACCESS_TOKEN = \"<PASTE_YOUR_TOKEN_HERE>\"\n",
    "\n",
    "modelpath=\"meta-llama/Meta-Llama-3-8B\"\n",
    "\n",
    "# Load model with gradient checkpointing enabled\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    modelpath,    \n",
    "    device_map=\"auto\",\n",
    "    torch_dtype=torch.bfloat16,\n",
    "    token=HUGGINGFACE_ACCESS_TOKEN,\n",
    ")\n",
    "model.config.use_cache = False\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(modelpath, token=HUGGINGFACE_ACCESS_TOKEN)\n",
    "tokenizer.pad_token = tokenizer.eos_token"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5cd73b7d",
   "metadata": {},
   "source": [
    "If you get a message similar to `WARNING:accelerate.big_modeling:Some parameters are on the meta device because they were offloaded to the cpu and disk.`, then the model itself is too big for your GPU. If you don't have a bigger / additional GPU at hand, you can use a quantization method like we show in the [QLoRA notebook](QLoRA_Llama_Finetuning.ipynb). Adding the quantization_config when loading the model and choosing a quantized `LoRAConfig` in the next step will enable quantized training."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "adapter_setup",
   "metadata": {},
   "source": [
    "## Initialize adapter\n",
    "\n",
    "We initialize the adapter functionality and add a LoRA adapter. When using gradient checkpointing with adapters, we need to enable input gradients explicitly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "init_adapter",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================================================================================\n",
      "Name                     Architecture         #Param      %Param  Active   Train\n",
      "--------------------------------------------------------------------------------\n",
      "lora_adapter             lora              3,407,872       0.085       1       1\n",
      "--------------------------------------------------------------------------------\n",
      "Full model                              4,015,263,744     100.000               0\n",
      "================================================================================\n"
     ]
    }
   ],
   "source": [
    "import adapters\n",
    "from adapters import LoRAConfig\n",
    "\n",
    "adapters.init(model)\n",
    "\n",
    "config = LoRAConfig()\n",
    "model.add_adapter(\"lora_adapter\", config=config)\n",
    "model.train_adapter(\"lora_adapter\")\n",
    "\n",
    "# Activate gradient checkpointing\n",
    "model.gradient_checkpointing_enable()\n",
    "\n",
    "print(model.adapter_summary())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "data_prep",
   "metadata": {},
   "source": [
    "## Prepare data for training\n",
    "\n",
    "The dataset is tokenized and truncated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "tokenize",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "def tokenize(element):\n",
    "    return tokenizer(\n",
    "        element[\"text\"],\n",
    "        truncation=True,\n",
    "        max_length=512,\n",
    "        add_special_tokens=False,\n",
    "    )\n",
    "\n",
    "dataset_tokenized = dataset.map(\n",
    "    tokenize, \n",
    "    batched=True, \n",
    "    num_proc=os.cpu_count(),\n",
    "    remove_columns=[\"text\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "training",
   "metadata": {},
   "source": [
    "## Training\n",
    "\n",
    "We specify training hyperparameters and train the model using the `AdapterTrainer` class. With gradient checkpointing enabled, we can use larger batch sizes than would otherwise be possible."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "training_args",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import TrainingArguments\n",
    "\n",
    "args = TrainingArguments(\n",
    "    output_dir=\"output/llama_gradient_checkpointing\",\n",
    "    per_device_train_batch_size=1,\n",
    "    per_device_eval_batch_size=1,\n",
    "    evaluation_strategy=\"steps\",\n",
    "    logging_steps=10,\n",
    "    save_steps=500,\n",
    "    eval_steps=187,\n",
    "    save_total_limit=3,\n",
    "    gradient_accumulation_steps=16,\n",
    "    max_steps=1875,\n",
    "    learning_rate=0.0002,\n",
    "    bf16=True,\n",
    "    warmup_ratio=0.03,\n",
    "    group_by_length=True,\n",
    "    lr_scheduler_type=\"constant\",\n",
    "    optim=\"adamw_torch\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "train",
   "metadata": {},
   "outputs": [],
   "source": [
    "from adapters import AdapterTrainer\n",
    "from transformers import DataCollatorForLanguageModeling\n",
    "\n",
    "trainer = AdapterTrainer(\n",
    "    model=model,\n",
    "    tokenizer=tokenizer,\n",
    "    data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),\n",
    "    train_dataset=dataset_tokenized[\"train\"],\n",
    "    eval_dataset=dataset_tokenized[\"test\"],\n",
    "    args=args,\n",
    ")\n",
    "\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "inference",
   "metadata": {},
   "source": [
    "## Inference\n",
    "\n",
    "For inference, we can disable gradient checkpointing since we don't need gradients:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "inference_setup",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Disable gradient checkpointing for inference\n",
    "model.gradient_checkpointing_disable()\n",
    "model.config.use_cache = True\n",
    "\n",
    "def prompt_model(model, text: str):\n",
    "    batch = tokenizer(f\"### Human: {text}\\n### Assistant:\", return_tensors=\"pt\")\n",
    "    batch = batch.to(model.device)\n",
    "    \n",
    "    model.eval()\n",
    "    with torch.inference_mode():\n",
    "        output_tokens = model.generate(**batch, max_new_tokens=50)\n",
    "\n",
    "    return tokenizer.decode(output_tokens[0], skip_special_tokens=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "test_inference",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(prompt_model(model, \"Explain gradient checkpointing in simple terms\"))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
